#!/usr/bin/env python3
"""
SAM模型调用接口

提供统一的SAM模型加载和推理接口，支持多种SAM实现
"""

import torch
import torch.nn as nn
import numpy as np
from typing import Dict, List, Tuple, Optional, Any, Union
from pathlib import Path
import sys

import sys
sys.path.append(str(Path(__file__).parent.parent.parent))
from utils.logger import get_logger


class SAMInterface:
    """
    SAM模型接口
    
    提供统一的SAM模型调用接口，支持：
    - SAM (Segment Anything Model)
    - SAM2
    - MedSAM2
    - MemSAM
    """
    
    def __init__(self,
                 model_type: str = 'sam_vit_h',
                 checkpoint_path: Optional[str] = None,
                 device: str = 'cuda',
                 sam_root: Optional[str] = None):
        """
        初始化SAM接口
        
        Args:
            model_type: 模型类型 ('sam_vit_h', 'sam_vit_b', 'sam_vit_l', 'sam2', 'medsam2', 'memsam')
            checkpoint_path: 检查点路径
            device: 设备 ('cuda' 或 'cpu')
            sam_root: SAM源码根目录（用于导入SAM模块）
                      如果为None，会尝试从默认位置加载：D:/mould/segment-anything-main
        """
        self.logger = get_logger("SAMInterface")
        self.model_type = model_type
        self.device = torch.device(device if torch.cuda.is_available() else 'cpu')
        
        # 如果没有提供sam_root，尝试默认路径
        if sam_root is None:
            default_sam_path = Path("D:/mould/segment-anything-main")
            if default_sam_path.exists():
                sam_root = str(default_sam_path)
                self.logger.info(f"使用默认SAM路径: {sam_root}")
        
        self.sam_root = sam_root
        self.model = None
        self.predictor = None
        
        # 加载模型
        self._load_model(checkpoint_path)
    
    def _load_model(self, checkpoint_path: Optional[str] = None):
        """加载SAM模型"""
        try:
            # 优先使用pip安装的segmentanything库
            try:
                from segment_anything import sam_model_registry, SamPredictor
                SAM_AVAILABLE = True
                self.logger.info("使用pip安装的segmentanything库")
            except ImportError:
                # 尝试从本地路径导入
                if self.sam_root:
                    sam_path = Path(self.sam_root)
                    if sam_path.exists():
                        sys.path.insert(0, str(sam_path))
                
                try:
                    from segment_anything import sam_model_registry, SamPredictor
                    SAM_AVAILABLE = True
                    self.logger.info("使用本地SAM路径")
                except ImportError:
                    # 尝试从external目录导入
                    try:
                        project_root = Path(__file__).parent.parent.parent.parent
                        external_sam = project_root / 'external' / 'segment-anything'
                        if external_sam.exists():
                            sys.path.insert(0, str(external_sam))
                            from segment_anything import sam_model_registry, SamPredictor
                            SAM_AVAILABLE = True
                            self.logger.info("使用external目录中的SAM")
                        else:
                            SAM_AVAILABLE = False
                    except ImportError:
                        SAM_AVAILABLE = False
            
            if not SAM_AVAILABLE:
                self.logger.warning("SAM模块未安装，将使用模拟实现")
                self.model = None
                self.predictor = None
                return
            
            if SAM_AVAILABLE:
                # 加载模型
                if checkpoint_path is None:
                    # 使用默认检查点路径
                    checkpoint_path = self._get_default_checkpoint()
                
                if checkpoint_path and Path(checkpoint_path).exists():
                    try:
                        self.model = sam_model_registry[self.model_type](checkpoint=checkpoint_path)
                        self.model.to(self.device)
                        self.predictor = SamPredictor(self.model)
                        self.logger.info(f"SAM模型加载成功: {self.model_type}")
                        self.logger.info(f"检查点路径: {checkpoint_path}")
                    except Exception as e:
                        self.logger.error(f"加载SAM模型时出错: {e}")
                        self.logger.warning("使用模拟实现")
                        self.model = None
                        self.predictor = None
                else:
                    self.logger.warning(f"检查点文件不存在: {checkpoint_path}")
                    self.logger.info("提示：请下载SAM检查点文件，或提供checkpoint_path参数")
                    self.logger.warning("使用模拟实现")
                    self.model = None
                    self.predictor = None
            else:
                self.logger.warning("segment_anything模块未找到")
                self.logger.info("请运行: pip install segmentanything")
                self.logger.warning("使用模拟实现")
                self.model = None
                self.predictor = None
                
        except Exception as e:
            self.logger.error(f"加载SAM模型失败: {e}")
            self.logger.warning("将使用模拟实现")
            self.model = None
            self.predictor = None
    
    def _get_default_checkpoint(self) -> Optional[str]:
        """获取默认检查点路径"""
        project_root = Path(__file__).parent.parent.parent.parent
        checkpoints = {
            'sam_vit_h': project_root / 'checkpoints' / 'sam' / 'sam_vit_h_4b8939.pth',
            'sam_vit_b': project_root / 'checkpoints' / 'sam' / 'sam_vit_b_01ec64.pth',
            'sam_vit_l': project_root / 'checkpoints' / 'sam' / 'sam_vit_l_0b3195.pth',
        }
        checkpoint_path = checkpoints.get(self.model_type)
        
        # 如果项目目录下不存在，尝试常见的位置
        if checkpoint_path and not checkpoint_path.exists():
            # 尝试用户主目录
            home_dir = Path.home()
            alternative_paths = [
                home_dir / '.cache' / 'sam' / checkpoint_path.name,
                home_dir / 'checkpoints' / checkpoint_path.name,
            ]
            for alt_path in alternative_paths:
                if alt_path.exists():
                    return str(alt_path)
        
        return str(checkpoint_path) if checkpoint_path else None
    
    def segment_with_points(self,
                           image: Union[np.ndarray, torch.Tensor],
                           points: np.ndarray,
                           point_labels: Optional[np.ndarray] = None) -> np.ndarray:
        """
        使用点提示进行分割
        
        Args:
            image: 输入图像 (H, W, 3) 或 (3, H, W)
            points: 点坐标 (N, 2) 格式为 [x, y]
            point_labels: 点标签 (N,) 1表示前景，0表示背景
            
        Returns:
            分割掩码 (H, W)
        """
        if self.predictor is None:
            return self._simulate_segmentation(image, 'points')
        
        # 转换图像格式
        if isinstance(image, torch.Tensor):
            image = image.cpu().numpy()
            if len(image.shape) == 3 and image.shape[0] == 3:
                image = image.transpose(1, 2, 0)
        
        # 确保图像是uint8格式
        if image.dtype != np.uint8:
            image = (image * 255).astype(np.uint8) if image.max() <= 1.0 else image.astype(np.uint8)
        
        # 设置图像
        self.predictor.set_image(image)
        
        # 准备点提示
        if point_labels is None:
            point_labels = np.ones(len(points), dtype=np.int32)
        
        # 执行分割
        masks, scores, logits = self.predictor.predict(
            point_coords=points,
            point_labels=point_labels,
            multimask_output=False
        )
        
        return masks[0]  # 返回最佳掩码
    
    def segment_with_box(self,
                         image: Union[np.ndarray, torch.Tensor],
                         box: np.ndarray) -> np.ndarray:
        """
        使用框提示进行分割
        
        Args:
            image: 输入图像 (H, W, 3) 或 (3, H, W)
            box: 边界框 (4,) 格式为 [x1, y1, x2, y2]
            
        Returns:
            分割掩码 (H, W)
        """
        if self.predictor is None:
            return self._simulate_segmentation(image, 'box')
        
        # 转换图像格式
        if isinstance(image, torch.Tensor):
            image = image.cpu().numpy()
            if len(image.shape) == 3 and image.shape[0] == 3:
                image = image.transpose(1, 2, 0)
        
        # 确保图像是uint8格式
        if image.dtype != np.uint8:
            image = (image * 255).astype(np.uint8) if image.max() <= 1.0 else image.astype(np.uint8)
        
        # 设置图像
        self.predictor.set_image(image)
        
        # 执行分割
        masks, scores, logits = self.predictor.predict(
            box=box,
            multimask_output=False
        )
        
        return masks[0]  # 返回最佳掩码
    
    def segment_with_mask(self,
                         image: Union[np.ndarray, torch.Tensor],
                         mask: np.ndarray) -> np.ndarray:
        """
        使用掩码提示进行分割
        
        Args:
            image: 输入图像 (H, W, 3) 或 (3, H, W)
            mask: 掩码提示 (H, W)
            
        Returns:
            分割掩码 (H, W)
        """
        if self.predictor is None:
            return self._simulate_segmentation(image, 'mask')
        
        # 转换图像格式
        if isinstance(image, torch.Tensor):
            image = image.cpu().numpy()
            if len(image.shape) == 3 and image.shape[0] == 3:
                image = image.transpose(1, 2, 0)
        
        # 确保图像是uint8格式
        if image.dtype != np.uint8:
            image = (image * 255).astype(np.uint8) if image.max() <= 1.0 else image.astype(np.uint8)
        
        # 设置图像
        self.predictor.set_image(image)
        
        # 执行分割
        masks, scores, logits = self.predictor.predict(
            mask_input=mask[None, :, :],
            multimask_output=False
        )
        
        return masks[0]  # 返回最佳掩码
    
    def segment(self,
               image: Union[np.ndarray, torch.Tensor],
               prompts: Dict[str, Any]) -> np.ndarray:
        """
        使用提示字典进行分割（统一接口）
        
        Args:
            image: 输入图像
            prompts: 提示字典，可能包含：
                - 'points': 点坐标
                - 'point_labels': 点标签
                - 'boxes': 边界框
                - 'masks': 掩码提示
                
        Returns:
            分割掩码
        """
        if 'points' in prompts:
            points = prompts['points']
            if isinstance(points, torch.Tensor):
                points = points.cpu().numpy()
            # 转换为 (N, 2) 格式
            if len(points.shape) == 3:
                points = points.reshape(-1, 2)
            point_labels = prompts.get('point_labels')
            if point_labels is not None and isinstance(point_labels, torch.Tensor):
                point_labels = point_labels.cpu().numpy()
            return self.segment_with_points(image, points, point_labels)
        
        elif 'boxes' in prompts:
            boxes = prompts['boxes']
            if isinstance(boxes, torch.Tensor):
                boxes = boxes.cpu().numpy()
            # 使用第一个框
            if len(boxes.shape) == 2:
                box = boxes[0]
            else:
                box = boxes
            return self.segment_with_box(image, box)
        
        elif 'masks' in prompts:
            mask = prompts['masks']
            if isinstance(mask, torch.Tensor):
                mask = mask.cpu().numpy()
            # 移除batch和时间维度
            if len(mask.shape) == 4:
                mask = mask[0, 0]
            elif len(mask.shape) == 3:
                mask = mask[0]
            return self.segment_with_mask(image, mask)
        
        else:
            raise ValueError("提示字典中必须包含 'points', 'boxes' 或 'masks' 之一")
    
    def _simulate_segmentation(self,
                              image: Union[np.ndarray, torch.Tensor],
                              prompt_type: str) -> np.ndarray:
        """
        模拟分割（当SAM未安装时使用）
        
        Args:
            image: 输入图像
            prompt_type: 提示类型
            
        Returns:
            模拟的分割掩码
        """
        # 获取图像尺寸
        if isinstance(image, torch.Tensor):
            if len(image.shape) == 4:
                _, _, h, w = image.shape
            elif len(image.shape) == 3:
                if image.shape[0] == 3:
                    h, w = image.shape[1], image.shape[2]
                else:
                    h, w = image.shape[:2]
            else:
                h, w = image.shape[:2]
        else:
            h, w = image.shape[:2]
        
        # 生成简单的模拟掩码（中心区域）
        mask = np.zeros((h, w), dtype=np.uint8)
        center_x, center_y = w // 2, h // 2
        radius = min(h, w) // 4
        y, x = np.ogrid[:h, :w]
        mask[(x - center_x)**2 + (y - center_y)**2 <= radius**2] = 1
        
        self.logger.warning(f"使用模拟分割结果（SAM未安装或未加载）")
        return mask
    
    def is_available(self) -> bool:
        """检查SAM是否可用"""
        return self.model is not None and self.predictor is not None

